[TRTLLM-14806][feat] Prefer Python V2 transceiver for LlamaForCausalLM and Gemma3ForCausalLM by moraxu · Pull Request #16787 · NVIDIA/TensorRT-LLM

moraxu · 2026-07-23T07:46:43Z

Dev Engineer Review

LlamaForCausalLM and Gemma3ForCausalLM now prefer the Python V2 transceiver when transceiver_runtime is left at auto and the effective backend is NIXL.
Explicit runtime selections are preserved, and UCX, MPI, or MOONCAKE still resolve to the C++ runtime.
The scope is limited to checkpoints implemented by these two classes; Llama4ForConditionalGeneration, MistralForCausalLM, and Gemma3ForConditionalGeneration are unchanged.
The upstream Llama/Gemma disaggregated tests, configs, and test-list registrations are retained.
Llama-3.1 and Gemma-3 accuracy tests select NIXL while omitting the runtime, exercising automatic V2 selection end to end.
Python V2 does not implement the serialized DataTransceiverState protocol and differs from C++ in experimental in-flight cancellation and synchronous timing output; those C++-specific tests remain on UCX/CPP, while the UCX logprobs RCCA path also remains unchanged.

QA Engineer Review

Unit coverage verifies that auto resolves to PYTHON for both model classes without overriding explicit runtime choices or selecting Python for unsupported backends.
The retained Llama-3.1 matrix covers overlap scheduling, block reuse, NGram/EAGLE3 speculative decoding, guided decoding, TP/PP topologies, and multiple instances on the automatically selected V2 path.
Gemma-3 coverage exercises automatic V2 selection with block reuse both enabled and disabled.
Legacy UCX/CPP tests remain scheduled for the behaviors that are specific to the C++ runtime, and the final diff removes no upstream test-list or waive entries.

Description

Prefer V2 transceiver backend for Gemma and Llama

Test Coverage

pytest -v tests/unittest/llmapi/test_llm_args.py::TestTransceiverRuntimeAutoResolution
LLM_MODELS_ROOT=/path/to/models pytest -v tests/integration/defs/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct
LLM_MODELS_ROOT=/path/to/models pytest -v tests/integration/defs/accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated.py::test_disaggregated_perf_metrics[TinyLlama-1.1B-Chat-v1.0]'
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0]'
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct]'
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer[False-TinyLlama-1.1B-Chat-v1.0]'
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer_missing_blocks[False-TinyLlama-1.1B-Chat-v1.0]'
LLM_MODELS_ROOT=/path/to/models pytest -v 'tests/integration/defs/disaggregated/test_disaggregated.py::test_llama4_long_context_kv_cache_overflow[llama4-models/nvidia/Llama-4-Maverick-17B-128E-Instruct-FP8]'

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

coderabbitai · 2026-07-23T07:54:32Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

Llama and Gemma causal language model classes now declare Python as their preferred transceiver runtime. Tests cover automatic resolution, while disaggregated tests explicitly select UCX or the C++ serialization protocol.

Changes

Transceiver runtime preferences

Layer / File(s)	Summary
Model runtime preference APIs `tensorrt_llm/_torch/models/modeling_gemma3.py`, `tensorrt_llm/_torch/models/modeling_llama.py`	Adds typed `get_preferred_transceiver_runtime` classmethods to both model classes, returning `"PYTHON"`.
Automatic runtime resolution coverage `tests/unittest/llmapi/test_llm_args.py`	Adds parameterized coverage confirming `"auto"` resolves to `"PYTHON"` for Llama and Gemma models.
Disaggregated protocol test configuration `tests/integration/defs/disaggregated/test_disaggregated.py`, `tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py`	Forces UCX for performance metrics tests and CPP for arbitrary KV-cache transfer tests.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#16524: Adds related model-level transceiver runtime preferences and automatic resolution coverage.

Suggested labels: api-compatible

Suggested reviewers: qijune, brnguyen2, litaotju

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the primary change: preferring the Python V2 transceiver for Llama and Gemma models.
Description check	✅ Passed	The description explains the change, scope, runtime behavior, preserved cases, and relevant unit and integration test coverage.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/modeling_gemma3.py (1)
279-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document both resolver-facing runtime-preference hooks.

Add matching Google-style docstrings documenting pretrained_config and the returned runtime sentinel:

tensorrt_llm/_torch/models/modeling_gemma3.py#L279-L284: document Gemma3ForCausalLM.get_preferred_transceiver_runtime.

tensorrt_llm/_torch/models/modeling_llama.py#L1137-L1142: document LlamaForCausalLM.get_preferred_transceiver_runtime.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_gemma3.py` around lines 279 - 284, Add
matching Google-style docstrings to
Gemma3ForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_gemma3.py lines 279-284 and
LlamaForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_llama.py lines 1137-1142, documenting the
pretrained_config parameter and the returned "PYTHON" runtime sentinel while
preserving the existing behavior.
Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma3.py`:
- Around line 279-284: Add matching Google-style docstrings to
Gemma3ForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_gemma3.py lines 279-284 and
LlamaForCausalLM.get_preferred_transceiver_runtime in
tensorrt_llm/_torch/models/modeling_llama.py lines 1137-1142, documenting the
pretrained_config parameter and the returned "PYTHON" runtime sentinel while
preserving the existing behavior.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 71444888-d1ad-47d4-8c6b-dcf66b86a774

📥 Commits

Reviewing files that changed from the base of the PR and between 93efd33 and 97231af.

📒 Files selected for processing (3)

tensorrt_llm/_torch/models/modeling_gemma3.py
tensorrt_llm/_torch/models/modeling_llama.py
tests/unittest/llmapi/test_llm_args.py

Shixiaowei02 · 2026-07-23T08:03:12Z

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

nv-xtf · 2026-07-23T08:32:35Z

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

I checked this, and the situation is slightly different from PR 16482. The existing Kimi test used DEFAULT, which caused the test harness to fall back to UCX/CPP, so it had to be explicitly changed to NIXL + PYTHON. Llama and Gemma already have corresponding NIXL + PYTHON disaggregated accuracy tests, so I don’t think similar test changes are needed here.
The existing test_auto_dtype cases will continue to exercise UCX/CPP and can be retained as legacy coverage.

Shixiaowei02 · 2026-07-23T13:02:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T13:09:09Z

PR_Github #61293 [ run ] triggered by Bot. Commit: 3b8aeaa Link to invocation

moraxu · 2026-07-23T17:02:05Z

Could you please check whether any existing tests need to be adjusted, similar to the changes made in PR 16482?

I only updated one test using TinyLlama in tests/integration/defs/disaggregated/test_disaggregated.py

moraxu · 2026-07-23T17:02:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T17:09:06Z

PR_Github #61343 [ run ] triggered by Bot. Commit: df03c3a Link to invocation

tensorrt-cicd · 2026-07-23T17:09:09Z

PR_Github #61293 [ run ] completed with state ABORTED. Commit: 3b8aeaa
/LLM/main/L0_MergeRequest_PR pipeline #49526 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

tensorrt-cicd · 2026-07-23T20:39:29Z

PR_Github #61343 [ run ] completed with state FAILURE. Commit: df03c3a
/LLM/main/L0_MergeRequest_PR pipeline #49570 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Shixiaowei02 · 2026-07-24T03:11:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T03:17:43Z

PR_Github #61484 [ run ] triggered by Bot. Commit: df03c3a Link to invocation

tensorrt-cicd · 2026-07-24T07:01:08Z

PR_Github #61484 [ run ] completed with state FAILURE. Commit: df03c3a
/LLM/main/L0_MergeRequest_PR pipeline #49704 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

tensorrt-cicd · 2026-08-07T22:57:58Z

PR_Github #64724 [ run ] triggered by Bot. Commit: 3beb801 Link to invocation

tensorrt-cicd · 2026-08-08T22:57:49Z

PR_Github #64724 [ run ] completed with state ABORTED. Commit: 3beb801

Link to invocation

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Shixiaowei02 · 2026-08-10T07:29:13Z

/bot run --disable-fail-fast

moraxu · 2026-08-10T22:27:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T22:35:20Z

PR_Github #65137 [ run ] triggered by Bot. Commit: d2d260d Link to invocation

tensorrt-cicd · 2026-08-11T04:56:19Z

PR_Github #65137 [ run ] completed with state FAILURE. Commit: d2d260d
/LLM/main/L0_MergeRequest_PR pipeline #52933 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

moraxu · 2026-08-11T05:37:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T05:43:23Z

PR_Github #65214 [ run ] triggered by Bot. Commit: d2d260d Link to invocation

tensorrt-cicd · 2026-08-11T12:25:00Z

PR_Github #65214 [ run ] completed with state FAILURE. Commit: d2d260d
/LLM/main/L0_MergeRequest_PR pipeline #52999 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

moraxu · 2026-08-11T16:25:37Z

/bot run

tensorrt-cicd · 2026-08-11T16:31:50Z

PR_Github #65333 [ run ] triggered by Bot. Commit: d2d260d Link to invocation

tensorrt-cicd · 2026-08-11T18:45:44Z

PR_Github #65333 [ run ] completed with state SUCCESS. Commit: d2d260d
/LLM/main/L0_MergeRequest_PR pipeline #53105 completed with status: 'SUCCESS'

CI Report

Link to invocation

…VIDIA#16787) Signed-off-by: Eric Tsai <ertsai@nvidia.com>

moraxu requested review from a team as code owners July 23, 2026 07:46

moraxu requested review from QiJune and lori-ren July 23, 2026 07:46

github-actions Bot assigned moraxu Jul 23, 2026

coderabbitai Bot reviewed Jul 23, 2026

View reviewed changes

Shixiaowei02 requested a review from nv-xtf July 23, 2026 08:03

nv-xtf approved these changes Jul 23, 2026

View reviewed changes

Shixiaowei02 approved these changes Jul 23, 2026

View reviewed changes

Shixiaowei02 force-pushed the dev-mguzek-use-transceiver-v2-for-llama-and-gemma branch from 97231af to 3b8aeaa Compare July 23, 2026 13:02

moraxu requested review from a team as code owners July 23, 2026 17:01

moraxu requested review from bo-nv, reasonsolo, ruodil and xinhe-nv July 23, 2026 17:01

coderabbitai Bot mentioned this pull request Jul 24, 2026

[None][feat] align time metrics between cpp and python cache transceiver #15871

Merged

1 task

moraxu added 10 commits August 10, 2026 15:29

First draft

5c61e81

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Final draft

8316427

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Fix

9743170

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Update tests

ed324aa

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Remove CPP tests for Llama

d4506b5

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Fix waives

f9222fb

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

precommit

a2a696a

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

[TRTLLM-14806][test] Restore disaggregated transceiver coverage

7381675

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

[TRTLLM-14806][test] Exercise automatic transceiver selection

c6df01e

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Update waives

d2d260d

Signed-off-by: Michal Guzek <mguzek@nvidia.com>

Shixiaowei02 force-pushed the dev-mguzek-use-transceiver-v2-for-llama-and-gemma branch from 3beb801 to d2d260d Compare August 10, 2026 07:29

moraxu merged commit a2f02da into NVIDIA:main Aug 11, 2026
7 checks passed

erictsai-nv added a commit to erictsai-nv/TensorRT-LLM that referenced this pull request Aug 12, 2026

[None][fix] Drop duplicate Gemma3 transceiver hook (added upstream in N…

5a1da50

…VIDIA#16787) Signed-off-by: Eric Tsai <ertsai@nvidia.com>

erictsai-nv mentioned this pull request Aug 12, 2026

[None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4 #17396

Merged

1 task

coderabbitai Bot mentioned this pull request Aug 12, 2026

[TRTLLM-14879][feat] qualify Qwen2 dense for MX #16974

Merged

erictsai-nv added a commit to erictsai-nv/TensorRT-LLM that referenced this pull request Aug 13, 2026

[None][fix] Drop duplicate Gemma3 transceiver hook (added upstream in N…

9023942

…VIDIA#16787) Signed-off-by: Eric Tsai <ertsai@nvidia.com>

erictsai-nv added a commit to erictsai-nv/TensorRT-LLM that referenced this pull request Aug 14, 2026

[None][fix] Drop duplicate Gemma3 transceiver hook (added upstream in N…

f9b4a48

…VIDIA#16787) Signed-off-by: Eric Tsai <ertsai@nvidia.com>

erictsai-nv mentioned this pull request Aug 28, 2026

[None][feat] Enable KVCacheManagerV2 by default for Llama and Llama4 #18342

Open

1 task

Uh oh!

Conversation

moraxu commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Shixiaowei02 commented Jul 23, 2026

Uh oh!

nv-xtf commented Jul 23, 2026

Uh oh!

Shixiaowei02 commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

moraxu commented Jul 23, 2026

Uh oh!

moraxu commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Shixiaowei02 commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 8, 2026

Uh oh!

Shixiaowei02 commented Aug 10, 2026

Uh oh!

moraxu commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

moraxu commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

moraxu commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants

moraxu commented Jul 23, 2026 •

edited

Loading

coderabbitai Bot commented Jul 23, 2026 •

edited

Loading